use std::collections::HashMap;

use syn::*;

#[derive(Debug, thiserror::Error)]
pub enum MirError {
    #[error("parse error: {0}")]
    Parse(String),
    #[error("expected a single free function in the file")] 
    ExpectedSingleFn,
    #[error("unsupported pattern in parameter or let binding")]
    UnsupportedPattern,
    #[error("unsupported statement")] 
    UnsupportedStmt,
    #[error("unsupported expression")] 
    UnsupportedExpr,
    #[error("if without else is not supported")] 
    IfWithoutElse,
}

pub type Result<T> = std::result::Result<T, MirError>;

/// Generate simplified MIR text from a Rust source string that contains exactly one function.
pub fn generate_mir(src: &str) -> Result<String> {
    let file = syn::parse_file(src).map_err(|e| MirError::Parse(e.to_string()))?;

    // Find the single free function
    let mut fns = file
        .items
        .into_iter()
        .filter_map(|it| match it {
            Item::Fn(f) => Some(f),
            _ => None,
        })
        .collect::<Vec<_>>();

    if fns.len() != 1 {
        return Err(MirError::ExpectedSingleFn);
    }
    let func = fns.remove(0);
    MirEmitter::new(func).emit()
}

struct MirEmitter {
    func: ItemFn,
    // variable/param name -> register index
    vars: HashMap<String, usize>,
    // next fresh register index (1-based; params occupy the first n)
    next_reg: usize,
    // ordered debug lines for parameters
    param_debugs: Vec<(String, usize)>,
    // textual output builder
    out: String,
    // current basic block buffer (lines inside the current bb)
    bb_lines: Vec<String>,
    // accumulated basic blocks (label -> lines)
    bbs: Vec<(usize, Vec<String>)>,
    // next basic block index to assign
    next_bb: usize,
    // When lowering inside an if-branch, suppress emitting `return;` lines;
    // branches should `goto -> bbJoin` and the final `return;` lives in join bb.
    suppress_return_stmt: bool,
}

impl MirEmitter {
    fn new(func: ItemFn) -> Self {
        Self {
            func,
            vars: HashMap::new(),
            next_reg: 1,
            param_debugs: Vec::new(),
            out: String::new(),
            bb_lines: Vec::new(),
            bbs: Vec::new(),
            next_bb: 0,
            suppress_return_stmt: false,
        }
    }

    fn emit(mut self) -> Result<String> {
        // Parameters -> registers _1.._n
        let (param_toks, ret_ty) = self.collect_signature()?;

        // Header: fn name(_1: T, ...) -> T {
        let name = self.func.sig.ident.to_string();
        self.push_line(0, &format!("fn {}({}) -> {} {{", name, param_toks, ret_ty));

        // Debug lines, if any
        for (name, reg) in self.param_debugs.clone() {
            self.push_line(1, &format!("debug {} => _{};", name, reg));
        }

        // Empty line between debug and first bb if there were any params
        if !self.param_debugs.is_empty() {
            self.out.push('\n');
        }

        // Start at bb0
        self.begin_bb();

        // Body statements
        let stmts = self.func.block.stmts.clone();
        for stmt in &stmts {
            self.lower_stmt(stmt)?;
        }

        // If no explicit return was emitted into a bb, ensure final return in current bb
        if !self.ends_with_return_in_any_bb() {
            // If return type is not unit, do nothing extra; spec examples all end with return;
            self.bb_lines.push(format!("{}return;", indent(2)));
        }

        // Close the last (or only) bb and write out all bbs
        self.end_current_bb();

        let bbs = std::mem::take(&mut self.bbs);
        for (idx, lines) in &bbs {
            // bbn header
            self.push_line(1, &format!("bb{}: {{", idx));
            for l in lines {
                self.out.push_str(l);
                self.out.push('\n');
            }
            self.push_line(1, "}");
            self.out.push('\n');
        }

        // Remove trailing blank line between last bb and closing brace, then close fn
        if self.out.ends_with('\n') {
            // remove the last newline we added after the final bb block
            let _ = self.out.pop();
        }
        self.push_line(0, "}");

        Ok(self.out)
    }

    fn collect_signature(&mut self) -> Result<(String, String)> {
        let mut parts = Vec::new();
        for (i, arg) in self.func.sig.inputs.iter().enumerate() {
            match arg {
                FnArg::Typed(PatType { pat, ty, .. }) => {
                    let name = pat_ident(pat)?;
                    let ty_str = ty_to_str(ty);
                    let reg = i + 1;
                    self.vars.insert(name.clone(), reg);
                    self.param_debugs.push((name.clone(), reg));
                    self.next_reg = reg + 1;
                    parts.push(format!("_{}: {}", reg, ty_str));
                }
                _ => return Err(MirError::UnsupportedPattern),
            }
        }

        let ret_ty = match &self.func.sig.output {
            ReturnType::Default => "()".to_string(),
            ReturnType::Type(_, ty) => ty_to_str(ty),
        };

        Ok((parts.join(", "), ret_ty))
    }

    fn lower_stmt(&mut self, stmt: &Stmt) -> Result<()> {
        match stmt {
            Stmt::Local(local) => self.lower_let(local),
            Stmt::Item(_) => Err(MirError::UnsupportedStmt),
            Stmt::Expr(expr, _) => self.lower_top_expr(expr),
            _ => Err(MirError::UnsupportedStmt),
        }
    }

    fn lower_top_expr(&mut self, expr: &Expr) -> Result<()> {
        match expr {
            Expr::Return(ExprReturn { expr: Some(val), .. }) => {
                // Peephole: if returning a just-bound let variable, and previous line is `_k = ...;`,
                // replace it with `_0 = ...;`.
                if let Expr::Path(ExprPath { path, .. }) = &**val {
                    if let Some(ident) = path.get_ident() {
                        if let Some(&k) = self.vars.get(&ident.to_string()) {
                            let prefix = format!("{}_{k} = ", indent(2));
                            if let Some(last) = self.bb_lines.last() {
                                if last.starts_with(&prefix) && last.trim_end().ends_with(';') {
                                    let mut last_line = self.bb_lines.pop().unwrap();
                                    // capture RHS inside the assignment
                                    let rhs = last_line[prefix.len()..].trim_end_matches(';').to_string();
                                    self.bb_lines
                                        .push(format!("{}{}_0 = {};", indent(2), maybe_indent(), rhs));
                                    if !self.suppress_return_stmt {
                                        self.bb_lines.push(format!("{}return;", indent(2)));
                                    }
                                    return Ok(());
                                }
                            }
                        }
                    }
                }

                // General case: lower into target register _0 directly
                self.lower_expr_into(val, Some("_0"))?;
                if !self.suppress_return_stmt {
                    self.bb_lines.push(format!("{}return;", indent(2)));
                }
                Ok(())
            }
            Expr::If(ExprIf { cond, then_branch, else_branch, .. }) => {
                let cond_reg = self.lower_bool_expr(cond)?; // ensures boolean op lowered

                let then_bb = self.next_bb_id();
                let else_bb = self.next_bb_id();
                let join_bb = self.next_bb_id();

                self.bb_lines.push(format!(
                    "{}switchInt({}) -> [1: bb{}, otherwise: bb{}];",
                    indent(2), cond_reg, then_bb, else_bb
                ));

                // Seal current bb
                self.end_current_bb();

                // THEN block
                self.begin_specific_bb(then_bb);
                let old = std::mem::replace(&mut self.suppress_return_stmt, true);
                for s in &then_branch.stmts {
                    self.lower_stmt(s)?;
                }
                self.suppress_return_stmt = old;
                // Ensure branch assigns _0 and goto -> join
                if !self.branch_assigned_return() {
                    // If not explicitly returned, fall back to goto join (without assigning _0)
                    // But spec examples always return in branches; still, emit a goto to join.
                }
                self.bb_lines
                    .push(format!("{}goto -> bb{};", indent(2), join_bb));
                self.end_current_bb();

                // ELSE block
                let else_block = else_branch
                    .as_ref()
                    .ok_or(MirError::IfWithoutElse)?;
                let else_block = match &*else_block.1 {
                    Expr::Block(ExprBlock { block, .. }) => block,
                    _ => return Err(MirError::UnsupportedExpr),
                };
                self.begin_specific_bb(else_bb);
                let old = std::mem::replace(&mut self.suppress_return_stmt, true);
                for s in &else_block.stmts {
                    self.lower_stmt(s)?;
                }
                self.suppress_return_stmt = old;
                if !self.branch_assigned_return() {}
                self.bb_lines
                    .push(format!("{}goto -> bb{};", indent(2), join_bb));
                self.end_current_bb();

                // JOIN block with final return
                self.begin_specific_bb(join_bb);
                self.bb_lines.push(format!("{}return;", indent(2)));
                // Do not end here; caller will flush later to preserve ordering
                Ok(())
            }
            _ => {
                // Top-level expression not explicitly supported
                Err(MirError::UnsupportedExpr)
            }
        }
    }

    fn lower_let(&mut self, local: &Local) -> Result<()> {
        // Only support: let ident = <expr>;
        let name = pat_ident(&local.pat)?;
        let expr = local
            .init
            .as_ref()
            .map(|init| &*init.expr)
            .ok_or(MirError::UnsupportedStmt)?;
        let val = self.lower_expr(expr)?;
        if let Some(idx) = parse_reg_index(&val) {
            // Expression already emitted into a register; bind variable to that register.
            self.vars.insert(name, idx);
        } else {
            // Literal or immediate; assign into a fresh register.
            let reg_idx = self.fresh_reg();
            let dest = format!("_{}", reg_idx);
            self.bb_lines
                .push(format!("{}{} = {};", indent(2), dest, val));
            self.vars.insert(name, reg_idx);
        }
        Ok(())
    }

    fn lower_expr(&mut self, expr: &Expr) -> Result<String> {
        self.lower_expr_into(expr, None)
    }

    fn lower_expr_into(&mut self, expr: &Expr, target: Option<&str>) -> Result<String> {
        match expr {
            Expr::Paren(ExprParen { expr, .. }) => self.lower_expr(expr),
            Expr::Lit(ExprLit { lit: Lit::Int(i), .. }) => {
                let lit = i.base10_digits().to_string();
                if let Some(dst) = target {
                    self.bb_lines
                        .push(format!("{}{} = {};", indent(2), dst, lit));
                    Ok(dst.to_string())
                } else {
                    Ok(lit)
                }
            }
            Expr::Path(ExprPath { path, .. }) => {
                let ident = path.get_ident().ok_or(MirError::UnsupportedExpr)?;
                let name = ident.to_string();
                if let Some(&idx) = self.vars.get(&name) {
                    let src = format!("_{}", idx);
                    if let Some(dst) = target {
                        self.bb_lines
                            .push(format!("{}{} = {};", indent(2), dst, src));
                        Ok(dst.to_string())
                    } else {
                        Ok(src)
                    }
                } else {
                    Err(MirError::UnsupportedExpr)
                }
            }
            Expr::Binary(ExprBinary { left, op, right, .. }) => {
                // arithmetic or boolean; decide op string
                let op_str = match op {
                    BinOp::Add(_) => "Add",
                    BinOp::Sub(_) => "Sub",
                    BinOp::Mul(_) => "Mul",
                    BinOp::Gt(_) => "Gt",
                    BinOp::Lt(_) => "Lt",
                    BinOp::Ge(_) => "Ge",
                    BinOp::Le(_) => "Le",
                    BinOp::Eq(_) => "Eq",
                    _ => return Err(MirError::UnsupportedExpr),
                };

                // Lower left then right to maintain left-to-right evaluation
                let l = self.lower_expr(left)?;
                let r = self.lower_expr(right)?;

                // Results of arithmetic/boolean ops are stored in a fresh register
                let dest = target
                    .map(|t| t.to_string())
                    .unwrap_or_else(|| format!("_{}", self.fresh_reg()));
                self.bb_lines.push(format!(
                    "{}{} = {}({}, {});",
                    indent(2), dest, op_str, l, r
                ));
                Ok(dest)
            }
            _ => Err(MirError::UnsupportedExpr),
        }
    }

    fn lower_bool_expr(&mut self, expr: &Expr) -> Result<String> {
        match expr {
            Expr::Binary(ExprBinary { op, .. }) if matches!(op, BinOp::Gt(_) | BinOp::Lt(_) | BinOp::Ge(_) | BinOp::Le(_) | BinOp::Eq(_)) =>
                self.lower_expr(expr),
            _ => Err(MirError::UnsupportedExpr),
        }
    }

    fn fresh_reg(&mut self) -> usize {
        let n = self.next_reg;
        self.next_reg += 1;
        n
    }

    fn begin_bb(&mut self) {
        let id = self.next_bb_id();
        self.begin_specific_bb(id);
    }

    fn begin_specific_bb(&mut self, idx: usize) {
        debug_assert!(self.bb_lines.is_empty());
        self.bbs.push((idx, Vec::new()));
        // We'll fill lines later; keep bb_lines bound to the last bb entry
    }

    fn end_current_bb(&mut self) {
        if let Some((_, lines)) = self.bbs.last_mut() {
            lines.extend(std::mem::take(&mut self.bb_lines));
        }
    }

    fn next_bb_id(&mut self) -> usize {
        let id = self.next_bb;
        self.next_bb += 1;
        id
    }

    fn ends_with_return_in_any_bb(&self) -> bool {
        self.bbs
            .iter()
            .any(|(_, lines)| lines.iter().any(|l| l.trim_end() == format!("{}return;", indent(2)).trim_end()))
            || self
                .bb_lines
                .iter()
                .any(|l| l.trim_end() == format!("{}return;", indent(2)).trim_end())
    }

    fn branch_assigned_return(&self) -> bool {
        self.bb_lines
            .iter()
            .any(|l| l.trim_start().starts_with("_0 = "))
    }

    fn push_line(&mut self, indent_lv: usize, s: &str) {
        self.out.push_str(&indent(indent_lv));
        self.out.push_str(s);
        self.out.push('\n');
    }
}

fn pat_ident(p: &Pat) -> Result<String> {
    match p {
        Pat::Ident(PatIdent { ident, .. }) => Ok(ident.to_string()),
        _ => Err(MirError::UnsupportedPattern),
    }
}

fn ty_to_str(ty: &Type) -> String {
    match ty {
        Type::Path(tp) => tp
            .path
            .segments
            .iter()
            .map(|s| s.ident.to_string())
            .collect::<Vec<_>>()
            .join("::"),
        _ => "()".to_string(),
    }
}

fn indent(level: usize) -> String { "    ".repeat(level) }

#[allow(dead_code)]
fn loc_str<T>(_node: &T) -> String {
    "unknown".to_string()
}

fn maybe_indent() -> &'static str { "" }

fn parse_reg_index(s: &str) -> Option<usize> {
    if let Some(rest) = s.strip_prefix('_') {
        rest.parse().ok()
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sample_4_no_params_lets_and_return() {
        let src = r#"
        pub fn experiment() -> i32 {
            let a = 1;
            let b = a + 5;
            let c = b * a;
            let d = c + 5;
            return d;
        }
        "#;
        let got = generate_mir(src).unwrap();
        let expect = r#"fn experiment() -> i32 {
    bb0: {
        _1 = 1;
        _2 = Add(_1, 5);
        _3 = Mul(_2, _1);
        _0 = Add(_3, 5);
        return;
    }
}"#;
        assert_eq!(got.trim(), expect.trim());
    }

    #[test]
    fn sample_1_if_else_params() {
        let src = r#"
        pub fn experiment(
            input1: i32,
            input2: i32,
        ) -> i32 {
            let a = input1 + input2 + 3;
            let b = input2 - 1;
            if b > 1 {
                return a + b;
            }
            else {
                return 0;
            }
        }
        "#;
        let got = generate_mir(src).unwrap();
        let expect = r#"fn experiment(_1: i32, _2: i32) -> i32 {
    debug input1 => _1;
    debug input2 => _2;

    bb0: {
        _3 = Add(_1, _2);
        _4 = Add(_3, 3);
        _5 = Sub(_2, 1);
        _6 = Gt(_5, 1);
        switchInt(_6) -> [1: bb1, otherwise: bb2];
    }

    bb1: {
        _0 = Add(_4, _5);
        goto -> bb3;
    }

    bb2: {
        _0 = 0;
        goto -> bb3;
    }

    bb3: {
        return;
    }
}"#;
        assert_eq!(got.trim(), expect.trim());
    }

    #[test]
    fn sample_10_many_params_add_chain() {
        let src = r#"
        pub fn experiment(
            input: i32, input2: i32, input3: i32, input4: i32,
            input5: i32, input6: i32,
        ) -> i32 {
            return input + input2 + input3 + input4 + input5 + input6;
        }
        "#;
        let got = generate_mir(src).unwrap();
        let expect = r#"fn experiment(_1: i32, _2: i32, _3: i32, _4: i32, _5: i32, _6: i32) -> i32 {
    debug input => _1;
    debug input2 => _2;
    debug input3 => _3;
    debug input4 => _4;
    debug input5 => _5;
    debug input6 => _6;

    bb0: {
        _7 = Add(_1, _2);
        _8 = Add(_7, _3);
        _9 = Add(_8, _4);
        _10 = Add(_9, _5);
        _0 = Add(_10, _6);
        return;
    }
}"#;
        assert_eq!(got.trim(), expect.trim());
    }
}
